# -*- coding: utf-8 -*-

from sympy import sympify, Pow, simplify, binomial, S
from sympy.abc import x, y
from mathsolver.functions.base.base import BaseValue


def get_mul_coeff_sym(poly, *num_sym):  # 获取乘式的系数和变量
    dt = poly.as_powers_dict()
    coeff = S.One
    sym = S.One
    for k, v in dt.items():
        if k.is_Number or k in num_sym:
            coeff *= k ** v
        else:
            sym *= k ** v
    return coeff, sym


class BaseErXiangShiValue(BaseValue):
    """向量变量"""

    def __init__(self, value):
        super(BaseErXiangShiValue, self).__init__(value)
        new_value = {}
        for key, item in value.items():
            new_value[key] = item
        self.value = new_value


class BaseErXiangShi(BaseValue):
    """二项式"""

    def __init__(self, value):
        super(BaseErXiangShi, self).__init__(value)
        self.name = "exs"
        self.flag = False  # False: 不能化简为二项式的标准式
        self.test = False
        poly = sympify(value)
        while not self.flag and not self.test:
            if isinstance(poly, Pow):
                di, mi = poly.args
                if len(di.args) == 2:
                    self.flag = True
                    self.a, self.b = simplify(di.args[0]), simplify(di.args[1])
                    self.n = mi
                    break
            self.test = True
            if not self.flag:
                poly = poly.factor()
        self.poly = poly
        self.num_sym = poly.free_symbols - {x, y}

    @staticmethod
    def instance(a, b, n):
        return BaseErXiangShi((a + b) ** n)

    # 通项公式
    def common(self):
        k = sympify("r")
        coef = binomial(self.n, k)
        tmp_a = self.a.factor()
        if tmp_a.is_Mul:  # 乘式
            coef_a, pl = get_mul_coeff_sym(tmp_a, *self.num_sym)  # 乘式的系数
            coef *= coef_a ** (self.n - k)
            tmp_a = pl

        if tmp_a.is_Pow:
            base, exp = tmp_a.args
            fir_a = base ** (exp * (self.n - k))
        else:
            fir_a = tmp_a ** (self.n - k)

        tmp_b = self.b.factor()
        if tmp_b.is_Mul:  # 乘式
            coef_b, pl = get_mul_coeff_sym(tmp_b, *self.num_sym)  # 乘式的系数
            coef *= coef_b ** k
            tmp_b = pl
        if tmp_b.is_Pow:
            base, exp = tmp_b.args
            fir_b = base ** (exp * k)
        else:
            fir_b = tmp_b ** k
        return BaseValue(coef * simplify(fir_a * fir_b))

    # 二项式
    def norm(self):
        return BaseValue((self.a + self.b) ** self.n)

    # 展开式
    def expand(self):
        return BaseValue(sympify((self.a + self.b) ** self.n).expand())

    # 第k+1项的值
    def the_rplus(self, t_r):
        return BaseValue(binomial(self.n, t_r) * sympify((self.a ** (self.n - t_r)) * (self.b ** t_r)))

    # 系数最大项的r
    def the_max_coeff(self):
        max_coeff = None
        max_index = []
        for i in range(self.n + 1):
            poly = self.the_rplus(i).value
            coeff = poly.as_coeff_Mul()[0]
            if not max_coeff:
                max_coeff = coeff
                max_index.append(i + 1)
            else:
                if max_coeff < coeff:
                    max_coeff = coeff
                    max_index = [i + 1]
                elif max_coeff == coeff:
                    max_index.append(i + 1)
        return max_index

    # 系数最小项的r
    def the_min_coeff(self):
        min_coeff = None
        min_index = []
        for i in range(self.n + 1):
            poly = self.the_rplus(i).value
            coeff = poly.as_coeff_Mul()[0]
            if not min_coeff:
                min_coeff = coeff
                min_index.append(i + 1)
            else:
                if min_coeff > coeff:
                    min_coeff = coeff
                    min_index = [i + 1]
                elif min_coeff == coeff:
                    min_index.append(i + 1)
        return min_index


if __name__ == '__main__':
    pass
